import rtree
import os
import cPickle 
from fasta import complement, Fasta
from genedex import Genedex
import numpy

class Pairdex(Genedex):

    def __setitem__(self, i, (feature_q, feature_s)):
        if i < len(self._store):
           del self[i]
           self._store.insert(i, (feature_q, feature_s))
        else:
            self._store.append((feature_q, feature_s))

        self.add(i, (feature_q['start'], feature_s['start'], feature_q['stop'], feature_s['stop']))

    # doesn't work?
    def __delitem__(self, i):
        feature_q, feature_s = self._store[i]
        self.delete(i, (feature_q['start'], feature_s['start'], feature_q['stop'], feature_s['start']))
        del self._store[i]

    def distance(self, a, b):
        raise Exception('not implemented')


    def __neighbors(self, f, n=1, strand=1, max_dist=60000, inclusive=False):
        raise Exception('not implemented')
        """ internal: only called when neighbors() is with
        a single strand, but no direction.
        so it calls neighbors() 2x, once for each direction, and picks
        out the best results.
        """
        n1 = self.neighbors(f, n=n + 1, strand=strand, max_dist=max_dist
                , direction=self.UP_STREAM, inclusive=False)
        # only use inclusive=True once, otherwise we get repeated
        # features.
        n2 = self.neighbors(f, n=n + 1, strand=strand, max_dist=max_dist
                , direction=self.DOWN_STREAM, inclusive=inclusive)

        if not isinstance(n1, list): n1 = [n1]
        if not isinstance(n2, list): n2 = [n2]
        neighbors = n1 + n2
        return self._cull_neighbors(f, neighbors, n)
    
    def neighbors(self, f, n=1, direction=None, strand=None, max_dist=60000, inclusive=False):
        raise Exception('not implemented')
        """\
        neighbors
        +++++++++
        The neighbors method returns the neighbors either up/downstream (or
        both) of a given feature. The arguments:
        ::

            f: a feature
            n: the number of neighbors to find.
            direction: -1 (upstream) or 1 (downstream), or left or right,
                        available as
                            *Genedex.UP_STREAM*
                            *Genedex.DOWN_STREAM*
                            *Genedex.LEFT*
                            *Genedex.RIGHT*
                    where upstream and downstream take into account the
                    strand of the feature, but LEFT and RIGHT do not.
                    LEFT will return features with lower basepair positions
            strand: to restrict the search to a single strand, use 1/-1 available as
                    *Genedex.PLUS*/*Genedex.MINUS*
                    default of None searches both strands
            max_dist: how far in basepairs to look before returning fewer than the
                        requested number of results. note, this can override
                        the n argument.
            inclusive: if False, start just outside the start/stop of f
                        if True, include the start/stop of f.

        The inclusive is useful when using an entry in the genedex instance as
        the feature but it is not wanted in the results (or vice-versa). By
        default, it wont appear in the results if only a single strand is
        requested. It will if no strand is specified (that may change in the
        future).
        """
        window = 3000
        if direction is None and strand is None and inclusive:
            return self.nearest(f, n)

        # they asked for a single strand, but no direction...
        if not direction in (Genedex.UP_STREAM, Genedex.DOWN_STREAM
                , Genedex.RIGHT, Genedex.LEFT):
            return self.__neighbors(f, n, strand, max_dist, inclusive)
        # strand == -1, direction == UP_STREAM   (-1) => add
        # strand == -1, direction == DOWN_STREAM ( 1) => subtract
        # strand ==  1, direction == DOWN_STREAM ( 1) => add
        # strand ==  1, direction == UP_STREAM   (-1) => subtract
        if direction in (Genedex.RIGHT, Genedex.LEFT):
            add_or_subtract = direction/abs(direction) # scale back to 1.
        else:
            add_or_subtract  = direction * f.get('strand', 1)
        neighbors = []
        # if it's (-), take start, (+) take stop.
        (astart, astop) = add_or_subtract == 1 and ('stop', 'start') or ('start', 'stop')
        start = f[astart] 
        if inclusive:
            # start at the other end, to get 
            # all overlapping features as well
            start = f[astop]
        else:
            # dont get this feature.
            start += add_or_subtract 

        # need to check we dont just continue on forever, so check the
        # distance from the original start point.
        start0 = start
        while len(neighbors) < n and abs(start0 - start) <= max_dist:
            stop = start + add_or_subtract * window
            # TODO: do this start < stop check in intersection?
            # spatialindex requires start < stop.
            istart, istop = stop < start and (stop, start) or (start, stop)
            new_neighbors = self.intersection({'start': istart, 'stop': istop}, strand=strand)
            if isinstance(new_neighbors, list):
                neighbors.extend(new_neighbors)
            else:
                neighbors.append(new_neighbors)
            start = stop + add_or_subtract


        if len(neighbors) <= n: return neighbors
        return self._cull_neighbors(f, neighbors, n)

    def _cull_neighbors(self, f,  neighbors, n):
        """ internal: 
        may have more than the desired # of neighbors because
        of the window. but cant just slice out the first n
        because we want to keep more than n if there are many that
        are equidistant.
        may return more than the requested number of features if there
        are many neighbors at equal distance to f.

        f: a feature
        neighbors: a list of features
        n: the number of features requested
        """
        def cmp_dist(a, b):
            return cmp(self.distance(a, f), self.distance(b, f))
        neighbors.sort(cmp=cmp_dist)

        if len(neighbors) <= n: return neighbors
        dist = self.distance(f, neighbors[n - 1])
        for i in range(n, len(neighbors)):
            if self.distance(f, neighbors[i]) > dist:
                return neighbors[:i]
        return neighbors

    def _load(self):
        """load the saved features from the pickle into _store"""
        self._store = cPickle.load(open(self.filename + '.pkl', 'rb'))

    def nearest(self, pair, n=1):
        """find the n nearest features to pair, where fq, fs = (pair) have 'start',
        'stop'. e.g.:
        strand not currently supported.

        >>> from genedex import Pairdex
        >>> g = Pairdex()
        >>> g.append({'start':1, 'stop':2})
        >>> g.append({'start':1000, 'stop':1001})
        >>> g.nearest({'start':22, 'stop':23})
        [{'start': 1, 'stop': 2}]
        >>> g.nearest({'start':22, 'stop':23}, n=2)
        [{'start': 1, 'stop': 2}, {'start': 1000, 'stop': 1001}]
        """
        # TODO (maybe) add strand
        f = pair
        return self.__getitem__(super(rtree.Rtree, self).nearest(\
                (f[0]['start'], f[1]['start'], f[0]['stop'], f[1]['stop']), n))


    def intersection(self, f):
        return self.__getitem__(super(rtree.Rtree, self).intersection(\
                (f[0]['start'], f[1]['start'], f[0]['stop'], f[1]['stop'])))


if __name__ == "__main__":
    import doctest
    doctest.testmod()
